// 

#include "TaskGenerator.h"

namespace ai_sim {

Define_Module(TaskGenerator);

void TaskGenerator::initialize() {
    tmp_tasktype = TaskType(static_cast<int>(par("tmp_tasktype")));
    // error("tasktype %d", tmp_tasktype);
    maxnumGpus = par("maxnumGpus");
    maxnumBatches = par("maxnumBatches");
    modelsize = par("modelsize");
    compute_time = par("compute_time");
    task_interval = par("task_interval");
    total_tasks = 0;
    omnetpp::cModule * controller_module =
        getParentModule()->getSubmodule("centralcontroller");
    controller = dynamic_cast<CentralController*> (controller_module);
    // TaskMsg *a_msg = new TaskMsg;
    // a_msg->setKind(TASK_ALLOC_MSG);
    // a_msg->setNow_task(gen_task());
    // send(a_msg, "central_controller_port");
    // controller->alloc_task(gen_fully_utilized_task());
    scheduleAt(omnetpp::simtime_t(0), new SelfMsg);
}

Task TaskGenerator::gen_fully_utilized_task() {
    int numGpus = maxnumGpus;
    int numBatches = maxnumBatches;
    Task tmp_task = {.taskId = total_tasks,
                    .tasktype = tmp_tasktype,
                    .numGpus = numGpus,
                    .numBatches = numBatches,
                    .modelsize = modelsize,
                    .compute_time = compute_time};
    total_tasks++;
    EV<< "A TASK GENERATED" << omnetpp::endl;
    return tmp_task;
}

void TaskGenerator::handleMessage(omnetpp::cMessage *msg) {
    if (msg->isSelfMessage() && total_tasks <1) 
    {
        // TaskMsg *a_msg = new TaskMsg;
        // a_msg->setKind(TASK_ALLOC_MSG);
        // a_msg->setNow_task(gen_task());
        // send(a_msg, "central_controller_port");
        controller->alloc_task(gen_fully_utilized_task());
        scheduleAt(omnetpp::simTime()+task_interval, msg);
    } else {
        delete msg;
    }
}
}  // namespace ai_sim
